#include <link.h>
#include <elf.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>

#include "linux_elf_utils.h"
#include "debug.h"

//------------------------------------------------------------------------------

TBGError ElfFindSegment(const int segment_type, unsigned long elf_load_addr,
                        Elf32_Phdr *ret_phdr) {
    TBGError ret = 0;
    CHECK_DBG_ERR(ret_phdr, NULL_PTR);
    
    Elf32_Ehdr ehdr = {};

    DEBUG("elf_load_addr = [%x]", elf_load_addr);

    // read elf header first
    CHECK_CALL(ReadByteData(elf_load_addr, &ehdr, sizeof(Elf32_Ehdr)),
               exit_point, ret);
    CHECK_RET(ehdr.e_ehsize, NO_ELF_HEADER, exit_point, ret);

    // e_phoff is the elh header offset
    addr_t phdr_addr = ELF_HEADER_BEGIN + ehdr.e_phoff; 
    CHECK_CALL(ReadByteData(phdr_addr, ret_phdr, sizeof(Elf32_Phdr)),
               exit_point, ret);
    CHECK_RET(ret_phdr->p_memsz, NO_PROGRAM_HEADER, exit_point, ret);

    // find the segment
    while(ret_phdr->p_type != segment_type && ret_phdr->p_type != PT_NULL) 
          CHECK_CALL(ReadByteData(phdr_addr += sizeof(Elf32_Phdr), ret_phdr,
                                  sizeof(Elf32_Phdr)),
                     exit_point, ret);
    
    CHECK_RET(ret_phdr->p_type != PT_NULL, NO_SUCH_SEGMENT, exit_point, ret);

exit_point:
    return ret;
}

//------------------------------------------------------------------------------

TBGError ElfLocateLinkmap(unsigned long elf_header_begin,
                          struct link_map *plink_map) {
    TBGError ret = 0;
    CHECK_DBG_ERR(plink_map, NULL_PTR);
    
	Elf32_Phdr phdr = {};
	Elf32_Dyn dyn = {}; 
	Elf32_Word got = 0;
    addr_t dyn_addr = 0, map_addr = 0;

    CHECK_CALL(ElfFindSegment(PT_DYNAMIC, elf_header_begin, &phdr),
               exit_point, ret); 
	
    // read the dynamic section header
	CHECK_CALL(ReadByteData(phdr.p_vaddr, &dyn, sizeof(Elf32_Dyn)),
               exit_point, ret);
	dyn_addr = phdr.p_vaddr;

    // find GOT section addr
	while(dyn.d_tag != DT_PLTGOT && dyn.d_tag != DT_NULL) 
        CHECK_CALL(ReadByteData(dyn_addr += sizeof(Elf32_Dyn), &dyn,
                                sizeof(Elf32_Dyn)), 
                   exit_point, ret);

    CHECK_RET(phdr.p_type != DT_NULL, NO_PLTGOT, exit_point, ret);

	got = (Elf32_Word) dyn.d_un.d_ptr;
    // link_map is the second GOT entry
	got += sizeof(Elf32_Word); 

    // get link_map addr 
	CHECK_CALL(ReadByteData((unsigned long) got, &map_addr, sizeof(Elf32_Word)),
               exit_point, ret);
    CHECK_RET(map_addr, INVALID_LINK_MAP_ADDR, exit_point, ret);
    // and read link_map. Ufffffff.
    CHECK_CALL(ReadByteData(map_addr, plink_map, sizeof(struct link_map)),
               exit_point, ret);
    CHECK_RET(plink_map->l_ld, INVALID_LINK_MAP, exit_point, ret);

exit_point:
    return ret;
}

//------------------------------------------------------------------------------

TBGError ElfGetDynSym(struct link_map *plink_map, Elf32_Sword type,
                      Elf32_Dyn *pdyn){
    TBGError ret = 0;
    CHECK_DBG_ERR(plink_map, NULL_PTR);
    CHECK_DBG_ERR(pdyn, NULL_PTR);

	unsigned long addr = (unsigned long) plink_map->l_ld;
    CHECK_DBG_ERR(addr, INVALID_LINK_MAP);

	CHECK_CALL(ReadByteData(addr, pdyn, sizeof(Elf32_Dyn)), exit_point, ret);

    while(pdyn->d_tag != DT_NULL && pdyn->d_tag != type) {
		addr += sizeof(Elf32_Dyn);
		CHECK_CALL(ReadByteData(addr, pdyn, sizeof(Elf32_Dyn)),
                   exit_point, ret);
	}

    if(pdyn->d_tag != type)
        ret = DYN_SYM_NOT_FOUND;

exit_point:
    return ret;
}

